import json
import os
import subprocess
import time
import uuid

from PIL import ImageFile
from dotenv import load_dotenv

import agent_config

from utils import sr_s3, sr_utils, sqs_video, status_codes
from utils.video_utils import check_content_review, process_video_with_watermark_service

StatusCodes = status_codes.StatusCodes

# 配置与初始化
load_dotenv(verbose=True, override=True)
ImageFile.LOAD_TRUNCATED_IMAGES = True  # 处理截断图片

# 环境变量与客户端初始化
app_env = os.getenv('APP_ENV', 'prod')

# 队列初始化 - 使用新的队列名
if app_env == 'dev':
    TASK_QUEUE = sqs_video.videoSQS.get_queue_by_name(QueueName='video_media_thumbmark_Test')
    result_queue_name = "AIHubVideoResult_Test"
else:
    TASK_QUEUE_NAME = agent_config.TASK_QUEUE_THUMBMART_NAME
    TASK_QUEUE = sqs_video.videoSQS.get_queue_by_name(QueueName=TASK_QUEUE_NAME)
    result_queue_name = "AIHubVideoResult"

print(TASK_QUEUE)


def send_result_to_queue(queue_name, body, status, msg, results=None, start_time=None, local_time=None, local_video=None, tokens=0, code=0, code_json=None):
    """发送结果到指定队列"""
    if local_video is None:
        local_video = {
            "width": 0,
            "height": 0,
            "fps": 0,
            "tokens": "0",
        }
    try:
        results = results or {}
        msg_body = {
            'QueueResult': queue_name,
            'status': status,
            'taskId': body.get('taskId'),
            'taskType': body.get('taskType'),
            'appId': body.get('appId'),
            'machineId': 1,
            'width': local_video['width'],
            'height': local_video['height'],
            'tokens': tokens,
            'fps': local_video['fps'],
            'results': results,
            'msgFetchEndTime': start_time,
            'start_time': start_time,
            'end_time': sr_utils.getTimes(local_time),
            'msg': msg,
            'code': code,
            'code_json': code_json
        }
        print("结果：", msg_body)
        queue = sqs_video.videoSQS.get_queue_by_name(QueueName=queue_name)
        queue.send_message(MessageBody=json.dumps(msg_body))
        print(f"已发送结果到队列 {queue_name}: {status}")
    except Exception as e:
        print(f"发送队列消息失败: {str(e)}")


# --------------------------
# 主循环逻辑
# --------------------------
def main():
    local_time = sr_utils.time_difference()
    print(local_time)
    while True:
        messages = TASK_QUEUE.receive_messages(MaxNumberOfMessages=1, WaitTimeSeconds=20)
        print(len(messages))
        if len(messages) == 0:
            time.sleep(1)
            continue
        # try:
        body = json.loads(messages[0].body)
        print(f"body: {body}")
        task_id = body.get('taskId', 0)
        if task_id == 0:
            sr_s3.set_logging(f'Error taskId is null: {task_id}')
            messages[0].delete()
            continue
        print(f"收到任务: {task_id}")

        # 清理
        sr_utils.delFiles('./download_vsr/')

        """统一处理视频生成任务（入口函数）"""
        start_time = sr_utils.getTimes(local_time)

        task_id = body.get('taskId', 'unknown')
        queue_result = result_queue_name

        # AIGC video
        results_attach_data = body.get('results_attach', {})
        resolution = results_attach_data.get('resolution', '480p')
        storage_service = results_attach_data.get('storageService', 's3')
        s3_prefix = results_attach_data.get('storedPrefix', 'video/image_to_video')
        bucket_name = results_attach_data.get('bucketName')
        origin_image = results_attach_data.get('origin_image')
        code_json = results_attach_data.get('code_json', {})
        results_attach = body.get('results_attach')

        if isinstance(results_attach, dict):
            tokens = results_attach.get('tokens')
        else:
            tokens = None  # 或者设置其他合理的默认值

        msgFetchStartTime = body.get('results_attach', {}).get('create_time', 0)

        # video 增强
        # 获取水印和缩略图控制参数
        watermark_enabled = body.get('results_attach', {}).get('watermark', '1') == '1'  # 默认开启水印
        thumbnail_enabled = body.get('results_attach', {}).get('thumbnail', '1') == '1'  # 默认开启缩略图

        msgFetchEndTime = sr_utils.getTimes(local_time)
        print('队列停留：{} 秒'.format((msgFetchEndTime - msgFetchStartTime) / 1000))

        if body.get('status', 1) == 0:
            err_msg = body.get("msg", "任务状态失败")
            send_result_to_queue(queue_result, body, 0, err_msg, start_time=start_time, local_time=local_time, code_json=code_json)
            messages[0].delete()
            continue

        local_video = {"width": body.get('width', 0), "height": body.get('height', 0), "fps": body.get('fps', 0), "tokens": tokens}
        # 检查 body 中是否存在 results 字段
        if not body.get('results'):
            err_msg = body.get("msg", "缺少参数results: results")
            send_result_to_queue(queue_result, body, 0, err_msg, start_time=start_time, local_time=local_time, code_json=code_json)
            messages[0].delete()
            continue

        url = body.get('results').get('key')
        # 必传参数校验
        if not url:
            err_msg = "缺少参数key: key"
            send_result_to_queue(queue_result, body, 0, err_msg, start_time=start_time, local_time=local_time, code_json=code_json)
            messages[0].delete()
            continue

        name = str(uuid.uuid4())
        donwload_path = './download_vsr/'
        donwload_img = './download_vsr/' + name + url[url.rfind('.'):]
        os.makedirs(donwload_path, exist_ok=True)
        fetchStartTime = sr_utils.getTimes(local_time)
        Res_s3, download_failed = sr_s3.download_file(url, donwload_img, serviceName=storage_service, bucketName=bucket_name)
        fetchEndTime = sr_utils.getTimes(local_time)
        donwload_time = (fetchEndTime - fetchStartTime) / 1000
        sr_s3.set_logging('donwload a file：{} seconds'.format(donwload_time))
        if Res_s3 is False:
            e = "download failed {}".format(str(download_failed))
            sr_s3.set_logging('download failed：{} '.format(body))
            results = {
                "origin_image": origin_image,
                "key": url,  # 使用超分后的url
                "watermark": "",
                "thumbnail": ""
            }
            error_msg = "下载超分视频失败【水印】"
            send_result_to_queue(queue_result, body, 1, error_msg, results=results, start_time=start_time, local_time=local_time, local_video=local_video, tokens=tokens, code="5011", code_json=code_json)
            messages[0].delete()
            continue

        # 调用指定参数的水印处理服务
        print(f"开始处理水印任务: {task_id}, 水印开启: {watermark_enabled}, 缩略图开启: {thumbnail_enabled}, 分辨率: {resolution}")
        sr_s3.set_logging(f"处理水印任务{task_id}: watermark={watermark_enabled}, thumbnail={thumbnail_enabled}, resolution={resolution}")

        # 调用指定参数的水印处理服务
        uploadStartTime = sr_utils.getTimes(local_time)
        watermark_result = process_video_with_watermark_service(
            donwload_img, task_id, resolution,
            storage_service=storage_service,
            bucket_name=bucket_name,
            s3_prefixes_path=s3_prefix,
            extract_thumbnail=thumbnail_enabled,  # 根据控制参数决定是否提取缩略图
            extract_watermark=watermark_enabled,  # 根据控制参数决定是否添加水印
        )
        uploadEndTime = sr_utils.getTimes(local_time)
        upload_time = (uploadEndTime - uploadStartTime) / 1000
        sr_s3.set_logging('Upload a file：{} seconds'.format(upload_time))

        print(watermark_result)

        # 处理结果并发送回队列
        if isinstance(watermark_result, dict) and watermark_result.get("status") == 1:
            # 成功处理
            results = {
                'msgFetchStartTime': msgFetchStartTime,
                'msgFetchEndTime': msgFetchEndTime,
                "origin_image": origin_image,
                "key": url,
                "watermark": watermark_result.get("watermark", ""),
                "thumbnail": watermark_result.get("thumbnail", ""),
                'fetchStartTime': fetchStartTime,
                'fetchEndTime': fetchEndTime,
                'uploadStartTime': uploadStartTime,
                'uploadEndTime': uploadEndTime
            }
            send_result_to_queue(queue_result, body, 1, "", results=results, start_time=start_time, local_time=local_time, local_video=local_video, tokens=tokens, code_json=code_json)
            sr_s3.set_logging(f"水印成功处理: {task_id}")
        else:
            # 处理失败
            results = {
                'msgFetchStartTime': msgFetchStartTime,
                'msgFetchEndTime': msgFetchEndTime,
                "origin_image": origin_image,
                "key": url,  # 使用超分后的url
                "watermark": "",
                "thumbnail": "",
                'fetchStartTime': fetchStartTime,
                'fetchEndTime': fetchEndTime,
                'uploadStartTime': uploadStartTime,
                'uploadEndTime': uploadEndTime,
            }
            error_msg = watermark_result.get("msg", "水印处理失败【水印】") if isinstance(watermark_result, dict) else "水印处理失败【水印】"
            send_result_to_queue(queue_result, body, 1, error_msg, results=results, start_time=start_time, local_time=local_time, local_video=local_video, tokens=tokens, code="5013", code_json=code_json)
            sr_s3.set_logging(f"水印处理失败: {task_id}")
        messages[0].delete()
        print(f"已删除队列消息: {task_id}")
        sr_s3.set_logging(f"已删除队列消息: {task_id}")
        # except Exception as e:
        #     print(f"主循环异常: {str(e)}")
        #     sr_s3.set_logging(f'Error: {str(e)}')
        #     messages[0].delete()
        #     time.sleep(1)


if __name__ == "__main__":
    main()
